package cn.wjee.gradle.flink.api;

import cn.wjee.gradle.flink.config.FlinkClusterClientBuilder;
import cn.wjee.gradle.flink.config.FlinkExtConfig;
import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.JobStatus;
import org.apache.flink.client.deployment.StandaloneClusterId;
import org.apache.flink.client.program.PackagedProgram;
import org.apache.flink.client.program.PackagedProgramUtils;
import org.apache.flink.client.program.rest.RestClusterClient;
import org.apache.flink.runtime.client.JobStatusMessage;
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.FileNotFoundException;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * Flink Submit Api
 *
 * @author listening
 */
public class FlinkSubmitterApi {
    private static final Logger log = LoggerFactory.getLogger(FlinkSubmitterApi.class);


    public static RestClusterClient<StandaloneClusterId> getFlinkClient(FlinkExtConfig config) {
        return FlinkClusterClientBuilder.builder()
                .defaultSsl()
                .address(config.getHost(), config.getRestPort(), config.getRpcPort())
                .build();
    }

    /**
     * Query Flink JOB ID By JobName
     *
     * @param jobName JobName
     * @param client  Config
     * @return List
     */
    public static List<JobID> getJobIdV2(RestClusterClient<StandaloneClusterId> client, String jobName) {
        try {
            final Collection<JobStatusMessage> jobStatusMessages = client.listJobs().get(15, TimeUnit.SECONDS);
            return jobStatusMessages.stream().filter(s -> s.getJobName().equals(jobName))
                    .filter(s -> s.getJobState() == JobStatus.RUNNING)
                    .map(JobStatusMessage::getJobId)
                    .collect(Collectors.toList());
        } catch (Exception e) {
            return new ArrayList<>();
        }
    }

    /**
     * Cancel Job
     *
     * @param restClusterClient Flink Client
     * @param jobId             JOB ID String
     * @param savePoint         SavePoint path
     */
    public static void cancelJob(RestClusterClient<StandaloneClusterId> restClusterClient, JobID jobId, String savePoint) {
        try {
            restClusterClient.cancelWithSavepoint(jobId, savePoint);
        } catch (Exception ignored) {
        }
    }

    /**
     * Submit Job To Flink
     *
     * @param config Config
     */
    public static void submitJob(FlinkExtConfig config) {
        RestClusterClient<StandaloneClusterId> flinkClient = null;
        PackagedProgram packagedProgram = null;
        try {
            File jobJar = Paths.get(config.getJobJar()).toAbsolutePath().toFile();
            if (!jobJar.exists()) {
                throw new FileNotFoundException("JarFile Not Exists");
            }
            flinkClient = getFlinkClient(config);

            packagedProgram = PackagedProgram.newBuilder()
                    .setJarFile(jobJar)
                    .setArguments(config.getJobArgs())
                    .build();

            JobGraph jobGraph = PackagedProgramUtils.createJobGraph(
                    packagedProgram,
                    flinkClient.getFlinkConfiguration(),
                    config.getParallelism(),
                    false
            );

            JobID existsJobId = getJobIdBeforeSubmit(flinkClient, config);
            if (existsJobId != null) {
                jobGraph.setJobID(existsJobId);
            }

            flinkClient.submitJob(jobGraph).get();
        } catch (Exception e) {
            log.error("submit Job fail", e);
        } finally {
            if (flinkClient != null) {
                flinkClient.close();
            }
            if (packagedProgram != null) {
                packagedProgram.close();
            }
        }
    }

    private static JobID getJobIdBeforeSubmit(RestClusterClient<StandaloneClusterId> flinkClient, FlinkExtConfig config) {
        try {
            final List<JobID> jobIdList = getJobIdV2(flinkClient, config.getJobName());
            if (jobIdList == null || jobIdList.isEmpty()) {
                return null;
            }
            jobIdList.forEach(s -> cancelJob(flinkClient, s, config.getSavePoint()));
            return jobIdList.stream().findFirst().orElse(null);
        } catch (Exception e) {
            return null;
        }
    }
}
